#!/usr/bin/python
# -*- coding: UTF-8 -*-
### This code is about the distributed formation control of consensus protocol with a certain
### Laplacian matrix and the formation transformation based on human recognization
### Code is made by Malan, email at malan9608@163.com
### Last modify on 18th Jun. 2021

#run : mpiexec -n 13 python formation.py

import rospy
from geometry_msgs.msg import Twist, Vector3, PoseStamped, TwistStamped
from std_msgs.msg import String 
import sys
import heapq
import numpy as np
import Queue
import time
import math
from pyquaternion import Quaternion
from mpi4py import MPI
from sensor_msgs.msg import Imu, NavSatFix
from mavros_msgs.msg import GlobalPositionTarget, State, PositionTarget
from geometry_msgs.msg import PoseStamped, Twist, TwistStamped
from mavros_msgs.srv import CommandBool, CommandTOL, SetMode, SetMavFrame
from naoji.msg import UAVDataMsg
from gazebo_msgs.srv import GetModelState
import copy

comm = MPI.COMM_WORLD
uav_id = comm.Get_rank()
uav_num = comm.Get_size()

class Formation:

    def __init__(self):

        self.hover = "HOVER"
        self.offboard = "OFFBOARD"
        self.id = uav_id
        self.uav_num = uav_num
        self.namespace = "iris_{proc}".format(proc=self.id)
        self.local_pose = PoseStamped()
        self.local_velocity = TwistStamped()
        self.cmd_vel_enu = Twist()
        self.arrive_print = True # determine whether the target position has been reached
        self.following_ids = [] # followers of this uav
        self.uav_data = UAVDataMsg()
        self.following_data = [UAVDataMsg() for i in range(self.uav_num)] # local position of other uavs, and only the position of followers of this uav is not zero
        self.follower_sub = [None]*self.uav_num
        self.arrive_count = 0
        self.loop_count = 0
        self.motion_type = 0  # 0:position 1:vel 2: accel
        self.frame_id = 7  # 1 :flu, 7: position
        self.received_new_task = False
        self.arm_state = False
        self.mavros_state = None
        self.received_imu = False
        self.frame = "BODY"
        self.mode = ""
        self.gcs_cmd = ""
        self.takeoff_target_pose = PoseStamped()
        self.target_vel = TwistStamped()

        # parameters
        self.Kp = 0.5
        self.Kpangz = 2
        self.velxy_max = 1  # 0.8
        self.velz_max = 0.5
        self.velangz_max = 2
        self.current_heading = 0.0
        self.f = 30  # control/communication rate
        # init height
        self.takeoff_height = 2
        self.target_yaw = 0.0
        self.bias = [[0,0] for i in range(uav_num)]
        self.origin_position = [0,-5,0]

        # ros publisher
        self.local_target_pub = rospy.Publisher(self.namespace + '/mavros/setpoint_raw/local', PositionTarget,
                                                queue_size=10)
        self.UAVdata_pub = rospy.Publisher(self.namespace + "/UAVdata",
                                                        UAVDataMsg, queue_size=10)
        # ros services
        self.armService = rospy.ServiceProxy(self.namespace + '/mavros/cmd/arming', CommandBool)
        self.takeoffService = rospy.ServiceProxy(self.namespace + '/mavros/cmd/takeoff', CommandTOL)
        self.landService = rospy.ServiceProxy(self.namespace + '/mavros/cmd/land', CommandTOL)
        self.flightModeService = rospy.ServiceProxy(self.namespace + '/mavros/set_mode', SetMode)
        self.frameService = rospy.ServiceProxy(self.namespace + '/mavros/setpoint_velocity/mav_frame', SetMavFrame)
        self.gazeboModelstate = rospy.ServiceProxy('gazebo/get_model_state', GetModelState)
        print(self.namespace, ": Px4 Controller Start!")

        # ros subscriber
        self.gcs_cmd_sub = rospy.Subscriber("/recognize_gcs/cmd", String, self.gcs_cmd_callback)
        self.mavros_sub = rospy.Subscriber(self.namespace + "/mavros/state", State, self.mavros_state_callback)
        self.imu_sub = rospy.Subscriber(self.namespace + "/mavros/imu/data", Imu, self.imu_callback)
        self.local_pose_sub = rospy.Subscriber(self.namespace + "/mavros/local_position/pose",
                                               PoseStamped, self.local_pose_callback)
        self.local_velocity_sub = rospy.Subscriber(self.namespace + "/mavros/local_position/velocity_local",
                                                   TwistStamped, self.local_velocity_callback)
        for i in range(self.uav_num):
            if i != self.id:
                self.follower_sub[i] = rospy.Subscriber("iris_"+str(i) + "/UAVdata",
                                                        UAVDataMsg, self.follower_callback, i)

    # callback functions
    def local_pose_callback(self, msg):
        self.local_pose = copy.deepcopy(msg)
        self.takeoff_target_pose = copy.deepcopy(msg)
        self.local_pose.pose.position.x += self.bias[self.id][0]
        self.local_pose.pose.position.y += self.bias[self.id][1]

    def follower_callback(self, msg, id):
        self.following_data[id] = copy.deepcopy(msg)

    def local_velocity_callback(self, msg):
        self.local_velocity = copy.deepcopy(msg)

    def mavros_state_callback(self, msg):
        self.mavros_state = msg.mode
        self.arm_state = msg.armed

    def imu_callback(self, msg):
        global global_imu, current_heading
        self.imu = msg
        self.current_heading = self.q2yaw(self.imu.orientation)
        self.received_imu = True

    def gcs_cmd_callback(self, msg):
        self.gcs_cmd = msg.data
        print(msg.data)
        self.arrive_print = True

    # main loop
    def loop(self):
        rospy.init_node(self.namespace + "_control_node")
        rate = rospy.Rate(self.f)
        distance = 10000
        # get init position of UAVs
        try:
            get_target_state = self.gazeboModelstate('iris_' + str(self.id), 'ground_plane')
            self.gazebo_target_pose = get_target_state.pose.position
            self.bias[self.id][0] = self.gazebo_target_pose.x
            self.bias[self.id][1]= self.gazebo_target_pose.y
        except rospy.ServiceException as e:
            print("Gazebo model state service" + self.id + "  call failed: %s") % e

        # loop
        while not rospy.is_shutdown():
            self.loop_count += 1
            # get information of myself
            self.construct_uav_data()
            self.UAVdata_pub.publish(self.uav_data)
            wait_time = 0
            # arrive?
            if self.arrive_count > 30 and self.arrive_print:
                print("Follower"+str(self.id)+":Arrived")
                self.arrive_print = False
                distance = 10000

            if not self.arrive_print:
                # Formation1:
                if self.mode == "AUTO.TAKEOFF" or self.mode == "formation1":
                    if self.mode != "formation1":
                        print("formation1")
                        target_position = self.takeoff_target_pose
                        self.read_set_file('UAV_pos1')
                    self.mode = "formation1"
                    self.following_ids = []
                    self.L_matrix = self.get_L_central_matrix()
                    # UAV0 is the leader and controlled by man (controlled by position)
                    if self.id == 0:
                        self.motion_type = 0
                        self.frame_id = 7
                        self.target_pose = self.construct_target(x=target_position.pose.position.x,
                                                                 y=target_position.pose.position.y,
                                                                 z=self.takeoff_height,
                                                                 yaw=self.target_yaw)
                        self.local_target_pub.publish(self.target_pose)
                    # other UAVs is followers and follow with the leader (controlled by consensus)
                    else:
                        self.motion_type = 1
                        self.frame_id = 1
                        for i in range(self.uav_num):
                            if self.L_matrix[self.id, i] == 1:
                                self.following_ids.append(i)
                        self.formation_control()
                        self.target_pose = self.construct_target(vx=self.target_vel.twist.linear.x,
                                                             vy=self.target_vel.twist.linear.y,
                                                             vz=self.target_vel.twist.linear.z,
                                                             yaw_rate=self.target_vel.twist.angular.z)
                        self.local_target_pub.publish(self.target_pose)
                # Land
                elif self.mode == "formation4":
                    self.mode = "AUTO.LAND"
                    if self.mavros_state != "AUTO.LAND":
                        self.flight_mode_set(mode='AUTO.LAND')
                    if (self.mavros_state == 'AUTO.LAND') and (self.local_pose.pose.position.z <= 0.2):
                        if self.arm_state == True:
                            self.disarm()

            # Take off
            if self.gcs_cmd == "FORM_1" and self.mode !="formation1":
                if self.mode != "AUTO.TAKEOFF":
                    target_position = self.takeoff_target_pose
                self.mode = "AUTO.TAKEOFF"
                self.motion_type = 0
                self.frame_id = 7
                if self.arm_state == False:
                    self.arm()
                if self.mavros_state != "OFFBOARD":
                    self.flight_mode_set(mode='OFFBOARD')
                if self.arrive_print:
                    distance = (self.takeoff_height - self.local_pose.pose.position.z) ** 2
                    self.target_pose = self.construct_target(x=target_position.pose.position.x,
                                                         y=target_position.pose.position.y,
                                                         z=self.takeoff_height,
                                                         yaw=self.current_heading)
                    self.local_target_pub.publish(self.target_pose)

            # Formation2:
            elif self.gcs_cmd == 'FORM_2':
                if self.mode != "formation2":
                    print("formation2")
                    target_position = self.takeoff_target_pose
                    self.read_set_file('UAV_pos1')
                self.mode = "formation2"
                self.following_ids = []
                self.L_matrix = self.get_L_central_matrix()
                if self.id == 0:
                    self.motion_type = 0
                    self.frame_id = 7
                    # go forward
                    target_position.pose.position.y += 0.02
                    self.target_pose = self.construct_target(x=target_position.pose.position.x,
                                                             y=target_position.pose.position.y,
                                                             z=self.uav_data.desired_pos.z,
                                                             yaw=self.target_yaw)
                    self.local_target_pub.publish(self.target_pose)

                else:
                    self.motion_type = 1
                    self.frame_id = 1
                    for i in range(self.uav_num):
                        if self.L_matrix[self.id, i] == 1:
                            self.following_ids.append(i)
                    self.formation_control()
                    self.target_pose = self.construct_target(vx=self.target_vel.twist.linear.x,
                                                     vy=self.target_vel.twist.linear.y,
                                                     vz=self.target_vel.twist.linear.z,
                                                     yaw_rate=self.target_vel.twist.angular.z)
                    self.local_target_pub.publish(self.target_pose)

            # Formation3
            elif self.gcs_cmd == 'FORM_3':
                if self.mode != "formation3":
                    print("formation3")
                    wait_time=0
                    target_position = self.takeoff_target_pose
                    self.read_set_file('UAV_pos2')
                self.mode = "formation3"
                self.following_ids = []
                self.L_matrix = self.get_L_central_matrix()
                if self.id == 0:
                    wait_time += 1
                    self.motion_type = 0
                    self.frame_id = 7
                    # go forward
                    target_position.pose.position.x -= 0.02
                    self.target_pose = self.construct_target(x=target_position.pose.position.x,
                                                             y=target_position.pose.position.y,
                                                             z=self.uav_data.desired_pos.z,
                                                             yaw=self.target_yaw)
                    self.local_target_pub.publish(self.target_pose)

                else:
                    self.motion_type = 1
                    self.frame_id = 1
                    for i in range(self.uav_num):
                        if self.L_matrix[self.id, i] == 1:
                            self.following_ids.append(i)
                    self.formation_control()
                    self.target_pose = self.construct_target(vx=self.target_vel.twist.linear.x,
                                                             vy=self.target_vel.twist.linear.y,
                                                             vz=self.target_vel.twist.linear.z,
                                                             yaw_rate=self.target_vel.twist.angular.z)
                    self.local_target_pub.publish(self.target_pose)

            #Formation4
            elif self.gcs_cmd == 'FORM_4':
                if self.mode != "formation4":
                    print("formation4")
                    wait_time = 0
                    target_position = self.takeoff_target_pose
                    self.read_set_file('UAV_pos3')
                self.mode = "formation4"
                self.following_ids = []
                self.L_matrix = self.get_L_central_matrix()
                if self.id == 0:
                    self.motion_type = 0
                    # go forward
                    if target_position.pose.position.y > 0.0:
                        target_position.pose.position.y -= 0.01
                    if target_position.pose.position.x < 0.0:
                        target_position.pose.position.x += 0.01
                    distance = (self.bias[self.id][0] - self.local_pose.pose.position.x) ** 2 + (self.bias[self.id][1] - self.local_pose.pose.position.y) ** 2
                    print(distance)
                    self.target_pose = self.construct_target(x=target_position.pose.position.x,
                                                             y=target_position.pose.position.y,
                                                             z=self.uav_data.desired_pos.z,
                                                             yaw=self.target_yaw)
                    self.local_target_pub.publish(self.target_pose)
                else:
                    self.motion_type = 1
                    self.frame_id = 1
                    for i in range(self.uav_num):
                        if self.L_matrix[self.id, i] == 1:
                            self.following_ids.append(i)
                    self.formation_control()
                    self.target_pose = self.construct_target(vx=self.target_vel.twist.linear.x,
                                                             vy=self.target_vel.twist.linear.y,
                                                             vz=self.target_vel.twist.linear.z,
                                                             yaw_rate=self.target_vel.twist.angular.z)
                    self.local_target_pub.publish(self.target_pose)
                    if self.local_pose.pose.position.z < 0.8:
                        self.arrive_count += 1
                    else:
                        self.arrive_count = 0

            # arrive_count
            if (self.mode == "AUTO.TAKEOFF" or self.mode == "formation4") and distance < 0.1:
                self.arrive_count += 1
            else:
                self.arrive_count = 0
            rate.sleep()

    def construct_uav_data(self):
        self.uav_data.header.stamp = rospy.Time.now()
        self.uav_data.uav_id = self.id
        self.uav_data.pose = self.local_pose.pose
        self.uav_data.velocity = self.local_velocity.twist

    # Laplacian matrix
    def get_L_central_matrix(self):
        L = np.zeros((self.uav_num,self.uav_num))
        for i in range(1,self.uav_num):
            L[i][0]=1
            L[i][i]=-1
        
        #L = np.array([[-0. ,0. , 0.,  0.,  0.,  0.],[ 1. ,-1. , 0. , 0. , 0. , 0.],[ 0.,  1. ,-1. , 0. , 0. , 0.],[ 0.,  0.,  1. ,-1. , 0. , 0.],[ 0. , 0. , 0. , 1. ,-1. , 0.],[ 0. , 0. , 0. , 0. , 1. ,-1.]])
        return L


    def formation_control(self):
        neighbor_num = len(self.following_ids)
        if len(self.following_ids) == 0:
            print("UAV ", self.id, "'s neighbot is zero!!!!!")
        neighbor_velocity_x = 0
        neighbor_velocity_y = 0
        neighbor_velocity_z = 0
        self.target_vel.twist.linear.x = 0
        self.target_vel.twist.linear.y = 0
        self.target_vel.twist.linear.z = 0
        self.target_vel.twist.angular.x = 0
        self.target_vel.twist.angular.y = 0
        self.target_vel.twist.angular.z = self.Kpangz * (self.target_yaw - self.current_heading)
        # print("neighbor_num",neighbor_num)
        for i in range(neighbor_num):
            self.target_vel.twist.linear.x += self.following_data[
                                                  i].pose.position.x - self.uav_data.pose.position.x - \
                                              self.following_data[i].desired_pos.x + \
                                              self.uav_data.desired_pos.x
            self.target_vel.twist.linear.y += self.following_data[
                                                  i].pose.position.y - self.uav_data.pose.position.y - \
                                              self.following_data[i].desired_pos.y + \
                                              self.uav_data.desired_pos.y
            self.target_vel.twist.linear.z += self.following_data[
                                                  i].pose.position.z - self.uav_data.pose.position.z - \
                                              self.following_data[i].desired_pos.z + \
                                              self.uav_data.desired_pos.z
            # eliminate error:
            neighbor_velocity_x += self.following_data[i].velocity.linear.x
            neighbor_velocity_y += self.following_data[i].velocity.linear.y
            neighbor_velocity_z += self.following_data[i].velocity.linear.z
        # print(self.target_vel.twist.linear)
        omega = self.Kp/neighbor_num
        self.target_vel.twist.linear.x = self.limit_amplitude(
            self.target_vel.twist.linear.x * omega + neighbor_velocity_x / neighbor_num, self.velxy_max)
        self.target_vel.twist.linear.y = self.limit_amplitude(
            self.target_vel.twist.linear.y * omega + neighbor_velocity_y / neighbor_num, self.velxy_max)
        self.target_vel.twist.linear.z = self.limit_amplitude(
            self.target_vel.twist.linear.z * omega + neighbor_velocity_z / neighbor_num, self.velz_max)
        self.target_vel.twist.angular.z = self.limit_amplitude(self.target_vel.twist.angular.z * self.Kpangz,
                                                               self.velangz_max)
        # self.trans_flu2enu()

    def read_set_file(self, txt_all_pos):
        all_pos_path='txt/'+txt_all_pos+'.txt'
        all_desired_position = np.loadtxt(all_pos_path)
        self.uav_data.desired_pos.x = all_desired_position[self.id][0]
        self.uav_data.desired_pos.y = all_desired_position[self.id][1]
        self.uav_data.desired_pos.z = all_desired_position[self.id][2]

    def construct_target(self, x=0, y=0, z=0, vx=0, vy=0, vz=0, afx=0, afy=0, afz=0, yaw=0, yaw_rate=0):
        target_raw_pose = PositionTarget()
        target_raw_pose.coordinate_frame = self.frame_id

        target_raw_pose.position.x = x
        target_raw_pose.position.y = y
        target_raw_pose.position.z = z

        target_raw_pose.velocity.x = vx
        target_raw_pose.velocity.y = vy
        target_raw_pose.velocity.z = vz

        target_raw_pose.acceleration_or_force.x = afx
        target_raw_pose.acceleration_or_force.y = afy
        target_raw_pose.acceleration_or_force.z = afz

        target_raw_pose.yaw = yaw
        target_raw_pose.yaw_rate = yaw_rate

        if (self.motion_type == 0):
            target_raw_pose.type_mask = PositionTarget.IGNORE_VX + PositionTarget.IGNORE_VY + PositionTarget.IGNORE_VZ \
                                        + PositionTarget.IGNORE_AFX + PositionTarget.IGNORE_AFY + PositionTarget.IGNORE_AFZ \
                                        + PositionTarget.IGNORE_YAW_RATE
        if (self.motion_type == 1):
            target_raw_pose.type_mask = PositionTarget.IGNORE_PX + PositionTarget.IGNORE_PY + PositionTarget.IGNORE_PZ \
                                        + PositionTarget.IGNORE_AFX + PositionTarget.IGNORE_AFY + PositionTarget.IGNORE_AFZ \
                                        + PositionTarget.IGNORE_YAW
        if (self.motion_type == 2):
            target_raw_pose.type_mask = PositionTarget.IGNORE_PX + PositionTarget.IGNORE_PY + PositionTarget.IGNORE_PZ \
                                        + PositionTarget.IGNORE_VX + PositionTarget.IGNORE_VY + PositionTarget.IGNORE_VZ \
                                        + PositionTarget.IGNORE_YAW

        return target_raw_pose

    def limit_amplitude(self, data, max_amplitude):
        if max_amplitude < 0:
            print('Warning! Max Amplitude should be positive!')
        if data <= -max_amplitude:
            data = -max_amplitude
        elif data >= max_amplitude:
            data = max_amplitude
        return data

    def q2yaw(self, q):
        if isinstance(q, Quaternion):
            rotate_z_rad = q.yaw_pitch_roll[0]
        else:
            q_ = Quaternion(q.w, q.x, q.y, q.z)
            rotate_z_rad = q_.yaw_pitch_roll[0]

        return rotate_z_rad

    def arm(self):
        if self.armService(True):
            return True
        else:
            print("Vehicle arming failed!")
            return False

    def takeoff(self):
        if self.takeoffService(True):
            return True
        else:
            print("Vehicle takeoff failed!")
            return False

    def disarm(self):
        if self.armService(False):
            return True
        else:
            print("Vehicle disarming failed!")
            return False

    def flight_mode_set(self, mode):
        """ mode selectable
        MANUAL, ACRO, ALTCTL, POSCTL, OFFBOARD, STABILIZED, RATTITUDE
        AUTO.MISSION, AUTO.LOITER, AUTO.RTL, AUTO.LAND, AUTO.RTGS, AUTO.READY, AUTO.TAKEOFF
        """
        if self.flightModeService(custom_mode=mode):
            return True
        else:
            print(self.namespace + mode + "Failed")

    def takeoff_detection(self):
        if self.local_pose.pose.position.z > 0.2 and self.arm_state:
            return True
        else:
            return False

if __name__ == '__main__':
    formation = Formation()
    formation.loop()
